# -*- coding: utf-8 -*-
"""
    PBAR_CompareCalibrationBasicScanZspec.py
# plot the 


Created on Wed May 07 17:04:25 2014

@author: jkwong
"""

# open the zspec data
import copy, glob, shutil, codecs, os, cPickle
import PBAR_Zspec, PBAR_Cargo
reload(PBAR_Zspec)
reload(PBAR_Cargo)
import numpy as np
import matplotlib.pyplot as plt
from scipy import interpolate
from matplotlib import cm
from scipy import ndimage

# Set useful plot variables
plotColors = ['r', 'b', 'g', 'm', 'c', 'y', 'k'] * 10
lineStyles = ['-', '-.', ':', '_', '|'] *  10
markerTypes = ['.', 'o', 'v', '^', '<', '>', '1', '2', '3', '4', 's', 'p', '*', 'h', 'H', '+', 'x', 'D', 'd']
markerTypes = markerTypes * 2
colorList = np.array(['r', 'b', 'g', 'm', 'c', 'y'])

# data input and output locations
dataPath = r'E:\PBAR\data4\BasicScansStandardWidth'
dataPath = r'E:\PBAR\data4\BasicScansStandardWidthPickle'
dataOutputBaseDir = r'E:\PBAR\data4\BasicScansStandardWidth'

# new compressed image width
compressedWidth = 400
compressedWidth2 = 1000

#################################
## DEFINE FILE NAMES

datStandardWidthList = []
datStandardWidthCompressedList = []
discrimStandardWidthCompressedFeaturesList = []

datCargoStandardWidthList = []
datCargoStandardWidthCompressedList = []

discrimStandardWidthCompressedFeaturesList2 = []

datCargoStandardWidthCompressedList2 = []

markerStandardWidthList = []

datasetDescription = PBAR_Zspec.ReadCargoDataDescriptionFile(r'E:\PBAR\data4\CargoSet2.txt')

energy = np.array([7.50E+00,1.05E+01,1.35E+01,1.65E+01,1.95E+01,2.25E+01,2.55E+01, \
    2.85E+01,3.15E+01,3.45E+01,3.75E+01,4.05E+01,4.35E+01,4.65E+01,4.95E+01,\
    5.25E+01,5.55E+01,5.85E+01,6.15E+01,6.45E+01,6.75E+01,7.05E+01,7.35E+01, \
    7.65E+01,7.95E+01,8.25E+01,8.55E+01,8.85E+01,9.15E+01,9.45E+01,9.75E+01,1.01E+02])

for datasetIndex, datasetDescript in enumerate(datasetDescription['scanID']):    
    print('Set %d' %datasetIndex)
#    if datasetIndex == 10:
#        break
    # open standard width files
    # Read Zspec
    filenameZspec = '%s-FDFC-All_SW.dat' %datasetDescription['scanID'][datasetIndex]
    fullFilenameZspec = os.path.join(dataPath, filenameZspec)
    print('Loading %s' %fullFilenameZspec)
    energy, datStandardWidth = PBAR_Zspec.ReadZspecBasicScanPickle(fullFilenameZspec)
    
    datStandardWidthList.append(datStandardWidth)
    datZspecSmall = PBAR_Zspec.ZspecBasicReduceSize(datStandardWidth, energy, compressedWidth)
    datStandardWidthCompressedList.append(datZspecSmall)
    # calculate the discriminant and features
    discrimSmall = PBAR_Zspec.CalculateFeaturesZspecBasic(datZspecSmall, energy, 1)
    discrimStandardWidthCompressedFeaturesList.append(discrimSmall)
    
    # do again at different size
    datZspecSmall = PBAR_Zspec.ZspecBasicReduceSize(datStandardWidth, energy, compressedWidth2)
    discrimSmall2 = PBAR_Zspec.CalculateFeaturesZspecBasic(datZspecSmall, energy, 1)
    discrimStandardWidthCompressedFeaturesList2.append(discrimSmall2)

    
    # Read Cargo
    filenameCargo = 'PBAR-%s.cargoimageSW.dat' %datasetDescription['dataFile'][datasetIndex]
    fullFilenameCargo = os.path.join(dataPath, filenameCargo)
    print('Loading %s' %fullFilenameCargo)
    with open(fullFilenameCargo, 'rb') as fid:
        datCargoStandardWidth = cPickle.load(fid)
    datCargoStandardWidthList.append(datCargoStandardWidth)
    
    datCargoStandardWidthCompressedList.append(PBAR_Zspec.CargoReduceSize(datCargoStandardWidth, datCargoStandardWidth.shape[1], compressedWidth))

    datCargoStandardWidthCompressedList2.append(PBAR_Zspec.CargoReduceSize(datCargoStandardWidth, datCargoStandardWidth.shape[1], compressedWidth2))

    # Read in marker files
    filenameMarker = filenameCargo.replace('cargoimageSW.dat', 'cargomarkerSW')
    fullFilenameMarker = fullFilenameCargo.replace('cargoimageSW.dat', 'cargomarkerSW')
    # some don't have marker files
    if os.path.exists(fullFilenameMarker):
    #        markers = PBAR_Cargo.ReadCargoMarker(fullFilenameMarker)
        markerStandardWidth = PBAR_Cargo.ReadCargoMarker(fullFilenameMarker)
    else:
        markerStandardWidth = []
    markerStandardWidthList.append(markerStandardWidth)



# Get the background points
discrimBackgroundList = []
featureNameList = ['binMean', 'binKur', 'binSkew', 'count', 'dist0']

cargoCountRange = [0.0, 4e7]
for datasetIndex in xrange(len(datStandardWidthList)):
    print(datasetIndex)
    discrimBackground = {}
    datStandardWidth = datStandardWidthList[datasetIndex]
    datCargoStandardWidth = datCargoStandardWidthList[datasetIndex]
    x = np.reshape(datCargoStandardWidth, datCargoStandardWidth.shape[0] * datCargoStandardWidth.shape[1])
    cut = (x > cargoCountRange[0]) & (x < cargoCountRange[1])
    # interp version
    discrim = PBAR_Zspec.CalculateFeaturesZspecBasic(datStandardWidth, energy, 1)
    for featureName in featureNameList:
        discrimBackground[featureName] = np.reshape(discrim[featureName], datCargoStandardWidth.shape[0] * datCargoStandardWidth.shape[1])[cut]
    discrimBackgroundList.append(discrimBackground)

basepath = r'C:\Users\jkwong\Documents\Work\PBAR\data4\Mar-files'
fullFilename = os.path.join(basepath, 'zspec3Set.dat')
with open(fullFilename ,'rb') as fid:
    print('reading %s' %fullFilename)
    datZspecCal =  cPickle.load(fid)


############
## PLOTS

# plot the features vs counts
# all detectors

# select the features
featureNameBasicScan = 'binMean'
featureNameCal = 'binMean_g1'
#
#featureNameBasicScan = 'binKur'
#featureNameCal = 'binKur'
#
featureNameBasicScan = 'binSkew'
featureNameCal = 'binSkew'

# get the feature index
featureIndex = np.where(datZspecCal['statsCollapseListALL'] == featureNameCal)[0][0]

i = 23


plt.figure()
plt.grid()
markerSize = 5

# plot the zspec data points

# no compress
plt.plot(discrimBackgroundList[i]['count'], discrimBackgroundList[i][featureNameBasicScan], \
    'sm',  alpha = 0.2, markersize = 3, label = 'Low Z from %s' %datasetDescription['scanID'][i])
#
#cut = (datCargoStandardWidthCompressedList[i] > cargoCountRange[0]) & (datCargoStandardWidthCompressedList[i] < cargoCountRange[1])
#plt.plot(discrimStandardWidthCompressedFeaturesList[i]['count'][cut], discrimStandardWidthCompressedFeaturesList[i][featureName][cut], \
#    'sm',  alpha = 0.2, markersize = 3, label = 'Low Z from %s' %datasetDescription['scanID'][i])
#
#cut = (datCargoStandardWidthCompressedList2[i] > cargoCountRange[0]) & (datCargoStandardWidthCompressedList2[i] < cargoCountRange[1])
#plt.plot(discrimStandardWidthCompressedFeaturesList2[i]['count'][cut], discrimStandardWidthCompressedFeaturesList2[i][featureName][cut], \
#    'sm',  alpha = 0.2, markersize = 3, label = 'Low Z from %s' %datasetDescription['scanID'][i])


for j in xrange(2):
    if j == 0:
        plotColor = 'r'
    else:
        plotColor = 'b'

    cut = datZspecCal['statsMatrix'][:,-1] == j
    plt.plot(datZspecCal['countMatrix'][cut], datZspecCal['statsMatrix'][cut,featureIndex], '.r', \
        color = plotColor, marker = markerTypes[i+1], markersize = markerSize,  alpha = 0.25, \
        label = 'Set, %s' %(datZspecCal['groupNamesExportList'][j]) )


datasetGroupsIndices = datZspecCal['datasetGroupsIndices']
groupNamesList = datZspecCal['groupNamesList']
stats = datZspecCal['stats']
detectorList = np.array([80])
# Make a dictionary of dictionaries containing a values from a subset of the detectors.
# statsCollapsed[<dataset group name>][<statistics name>] = 
#        values for all detectors specified for all datasets in the group
#

markerSize = 15
statsCollapsed = PBAR_Zspec.CreateStatsCollapsed(stats, groupNamesList, datasetGroupsIndices, detectorList)

plt.plot(statsCollapsed['PbNOT_0']['count'], statsCollapsed['PbNOT_0'][featureNameCal], 'sr', \
    markersize = markerSize,  alpha = 0.75, \
    label = 'Detector %d, %s' %(detectorList[0], 'PbNOT') )

plt.plot(statsCollapsed['PbALL']['count'], statsCollapsed['PbALL'][featureNameCal], 'sb', \
    markersize = markerSize,  alpha = 0.75, \
    label = 'Detector %d, %s' %(detectorList[0], 'PbALL') )

plt.legend(prop={'size':16})
plt.xlabel('Counts', fontsize = 16)
plt.ylabel(featureNameBasicScan, fontsize = 16)

if featureNameBasicScan == 'binMean':
    plt.ylim((0, 100))
elif featureNameBasicScan == 'binSkew':
    plt.ylim((-0.5, 3))
plt.xlim((0, 400))



# 
# show the pixels in the image that have a high skew value



# compare
i = 23

featureName = 'binSkew'

# cut in cargoimage count
cut1 = (datCargoStandardWidthCompressedList2[i] > cargoCountRange[0]) & (datCargoStandardWidthCompressedList2[i] < cargoCountRange[1])
# cut in the feature
cut2 = discrimStandardWidthCompressedFeaturesList2[i][featureName] > 1.75

temp = copy.copy(discrimStandardWidthCompressedFeaturesList2[i]['count'])

# show whole image
plt.figure()
plt.grid()
plt.imshow(datCargoStandardWidthCompressedList2[i].T, interpolation = 'nearest', aspect='auto', cmap = cm.Greys_r)
plt.xlim((0, 1000))
plt.title('cargoimage')


# with cargoimage mask 
plt.figure()
plt.grid()
plt.imshow(datCargoStandardWidthCompressedList2[i].T * cut1.T, interpolation = 'nearest', aspect='auto', cmap = cm.Greys_r)
plt.xlim((0, 1000))
plt.title('carogimage with count cut')

# with cargoimage mask and feature mask
plt.figure()
plt.grid()
plt.imshow(datCargoStandardWidthCompressedList2[i].T * cut1.T * cut2.T, interpolation = 'nearest', aspect='auto', cmap = cm.Greys_r)
plt.xlim((0, 1000))
plt.title('carogimage with count cut and feature cut')




# compare
plt.figure()
plt.grid()

plt.imshow(np.log(datZspecSmall.sum(2).T), interpolation = 'nearest', aspect='auto', cmap = cm.Greys_r)
xlim((0, 400))

plt.figure()
plt.grid()

plt.imshow(np.log(datStandardWidth.sum(2).T), interpolation = 'nearest', aspect='auto', cmap = cm.Greys_r)
xlim((0, 2000))


plt.figure()
plt.grid()

plt.imshow(datCargoStandardWidth.T, interpolation = 'nearest', aspect='auto', cmap = cm.Greys_r)
xlim((0, 2000))


plt.figure()
plt.grid()

plt.imshow(datCargoSmall.T, interpolation = 'nearest', aspect='auto', cmap = cm.Greys_r)
xlim((0, 400))


i = 1
plt.figure()
plt.grid()

plt.imshow(discrim[featureName].T, interpolation = 'nearest', aspect='auto', cmap = cm.Greys_r)
colorbar()
xlim((0, 2000))

plt.figure()
plt.grid()

plt.imshow(discrimSmall[featureName].T, interpolation = 'nearest', aspect='auto', cmap = cm.Greys_r)
colorbar()

xlim((0, 400))




i = 1
plt.figure()
plt.grid()

featureName = 'binSkew'

x = np.linspace(0,100, discrim[featureName].shape[0])
plt.plot(x, discrim[featureName].mean(1), '-k')
x = np.linspace(0,100, discrimSmall[featureName].shape[0])

plt.plot(x, discrimSmall[featureName].mean(1), '-r')

plt.xlim((0, 100))






# Old stuff for validating the compressed images



# Plot of zspec count versus a feature
# - plot with no carogimage count cut and one with a cut
# - shows that the carogimage cut removes a lot of bad data points

featureName = 'binMean'
featureName = 'binSkew'
plt.figure()

cut = ~np.isnan(discrimSmall[featureName]) & (datCargoSmall > cargoCountRange[0]) & (datCargoSmall < cargoCountRange[1])
plt.plot(discrimSmall['count'].flatten(), discrimSmall[featureName].flatten(), '.k')
plt.plot(discrimSmall['count'][cut].flatten(), discrimSmall[featureName][cut].flatten(), '.r')


plt.figure()
cut = ~np.isnan(discrim[featureName]) & (datCargoStandardWidth > cargoCountRange[0]) & (datCargoStandardWidth < cargoCountRange[1])
plt.plot(discrim['count'].flatten(), discrim[featureName].flatten(), '.k')
plt.plot(discrim['count'][cut], discrim[featureName][cut], '.r')


# Plot of the features vs zspec count for standard width and compressed images
#   - compares features calculated from standard width and compressed


#featureName = 'binMean'
featureName = 'binSkew'
plt.figure()
plt.grid()

# standwidth width
for i in xrange(5):
    plt.plot(discrimBackgroundList[i]['count'], discrimBackgroundList[i][featureName], '.r', color = plotColors[i],  alpha = 0.1, label = 'Standard Width, %d' %i)

## compressed
#for i in xrange(5):
#    cut =  (datCargoStandardWidthCompressedList[i] > cargoCountRange[0]) & (datCargoStandardWidthCompressedList[i] < cargoCountRange[1])
#    plt.plot(discrimStandardWidthCompressedFeaturesList[i]['count'][cut], discrimStandardWidthCompressedFeaturesList[i][featureName][cut], \
#        'sk', color = plotColors[i],  alpha = 0.2, markersize = 10, label = 'Compressed, %d' %i)

plt.legend()
plt.xlabel('Zspec Counts', fontsize = 16)
plt.ylabel(featureName, fontsize = 16)



# Plot of the features vs count for the background points
# with cargoimage counts instead of zspec counts
#featureName = 'binMean'
featureName = 'binSkew'
featureName = 'binSkew'

plt.figure()
plt.grid()


# standwidth width
for i in xrange(3):
#    plt.plot(discrimBackgroundList[i]['count'], di2scrimBackgroundList[i][featureName], '.r', color = plotColors[i],  alpha = 0.2, label = 'Standard Width, %d' %i)
    cut = (datCargoStandardWidthList[i] > cargoCountRange[0]) & (datCargoStandardWidthList[i] < cargoCountRange[1])
    plt.plot(datCargoStandardWidthList[i][cut].flatten(), discrimBackgroundList[i][featureName], '.r', color = plotColors[i],  alpha = 0.1, label = 'Standard Width, %d' %i)

# compressed
for i in xrange(3):
    cut =  (datCargoStandardWidthCompressedList[i] > cargoCountRange[0]) & (datCargoStandardWidthCompressedList[i] < cargoCountRange[1])
    plt.plot(datCargoStandardWidthCompressedList[i][cut], discrimStandardWidthCompressedFeaturesList[i][featureName][cut], \
        'sk', color = plotColors[i],  alpha = 0.2, markersize = 10, label = 'Compressed, %d' %i)

plt.legend()
plt.xlabel('CargoImage Counts', fontsize = 16)
plt.ylabel(featureName, fontsize = 16)






# compare
plt.figure()
plt.grid()

plt.imshow(np.log(datZspecSmall.sum(2).T), interpolation = 'nearest', aspect='auto', cmap = cm.Greys_r)
xlim((0, 400))

plt.figure()
plt.grid()

plt.imshow(np.log(datStandardWidth.sum(2).T), interpolation = 'nearest', aspect='auto', cmap = cm.Greys_r)
xlim((0, 2000))


plt.figure()
plt.grid()

plt.imshow(datCargoStandardWidth.T, interpolation = 'nearest', aspect='auto', cmap = cm.Greys_r)
xlim((0, 2000))


plt.figure()
plt.grid()

plt.imshow(datCargoSmall.T, interpolation = 'nearest', aspect='auto', cmap = cm.Greys_r)
xlim((0, 400))


i = 1
plt.figure()
plt.grid()

plt.imshow(discrim[featureName].T, interpolation = 'nearest', aspect='auto', cmap = cm.Greys_r)
colorbar()
xlim((0, 2000))

plt.figure()
plt.grid()

plt.imshow(discrimSmall[featureName].T, interpolation = 'nearest', aspect='auto', cmap = cm.Greys_r)
colorbar()

xlim((0, 400))

i = 1
plt.figure()
plt.grid()

featureName = 'binSkew'

x = np.linspace(0,100, discrim[featureName].shape[0])
plt.plot(x, discrim[featureName].mean(1), '-k')
x = np.linspace(0,100, discrimSmall[featureName].shape[0])

plt.plot(x, discrimSmall[featureName].mean(1), '-r')

plt.xlim((0, 100))


